package com.adnaan.backend.config;

import com.adnaan.backend.service.AiService;
import com.adnaan.backend.websocket.AiChatWebSocketHandler;
import com.adnaan.backend.websocket.SystemMessageWebSocketHandler;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    private final AiService aiService;

    public WebSocketConfig(AiService aiService) {
        this.aiService = aiService;
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(aiChatWebSocketHandler(), "/ai/ws")
                .setAllowedOrigins("*");
        registry.addHandler(systemMessageWebSocketHandler(), "/message/ws")
                .setAllowedOrigins("*");
    }

    @Bean
    public ObjectMapper objectMapper() {
        ObjectMapper objectMapper = new ObjectMapper();
        objectMapper.registerModule(new JavaTimeModule());
        return objectMapper;
    }

    @Bean
    public AiChatWebSocketHandler aiChatWebSocketHandler() {
        return new AiChatWebSocketHandler(aiService);
    }

    @Bean
    public SystemMessageWebSocketHandler systemMessageWebSocketHandler() {
        return new SystemMessageWebSocketHandler(objectMapper());
    }
} 